//
// Created by mio on 25-5-14.
//

#include "RtmpPushManager.h"

#include <QGuiApplication>
#include <QScreen>

#include "DXGICapture.h"
#include "AudioEncoder.h"
#include "VideoEncoder.h"
#include "AudioCapture.h"
#include "H264Paraser.h"

RtmpPushManager::RtmpPushManager() : audio_encoder_(nullptr), video_encoder_(nullptr), pusher_(nullptr),
        audio_Capture_(nullptr), screen_Capture_(nullptr)
{
    loop_ = new EventLoop(1);
}

RtmpPushManager::~RtmpPushManager()
{

}

bool RtmpPushManager::Open(const QString &str)
{
    if (!Init()) return false;
    if (pusher_->OpenUrl(str.toStdString(), 1000) < 0) return false;

    isConnect = true;

    videoCaptureThread_.reset(new std::thread([this]() {
        this->EncodeVideo();
    }));

    audioCaptureThread_.reset(new std::thread([this]() {
        this->EncodeAudio();
    }));
    qDebug() << "RtmpPushManager::Open";
    return true;
}

bool RtmpPushManager::Init()
{
    pusher_ = RtmpPublisher::Create(loop_);
    pusher_->SetChunkSize(60000);

    screen_Capture_.reset(new DXGICapture());
    if (!screen_Capture_->Init()) return false;

    QScreen *screen = QGuiApplication::primaryScreen();
    QRect screenGeometry = screen->geometry();
    qreal dpr = screen->devicePixelRatio();
    QSize physicalSize = QSize(screenGeometry.width() * dpr, screenGeometry.height() * dpr);

    video_encoder_.reset(new VideoEncoder());
    if (!video_encoder_->Open(physicalSize.width(), physicalSize.height(), 30, 80000, AV_PIX_FMT_BGRA)) return false;

    audio_Capture_.reset(new AudioCapture());
    if (!audio_Capture_->Init()) return false;

    audio_encoder_.reset(new AudioEncoder());
    if (!audio_encoder_->Open(audio_Capture_->GetSampleRate(), audio_Capture_->GetChannels(), AV_SAMPLE_FMT_S16, 64)) return false;

    //获取音频视频编码参数
    MediaInfo mediaInfo;
    uint8_t extradata[1024] = { 0 };
    int extradataSize = 0;

    /* 获取H264编码参数 */
    extradataSize = video_encoder_->GetSequenceParams(extradata, 1024);
    if (extradataSize <= 0) return false;

    //获取sps pps
    H264Paraser::Nal sps = H264Paraser::findNal(extradata, extradataSize);
    if (sps.first != nullptr && sps.second != nullptr && (*sps.first & 0x1f) == 7) // sps数据
    {
        mediaInfo.sps_size = sps.second - sps.first + 1;
        mediaInfo.sps.reset(new uint8_t[mediaInfo.sps_size], std::default_delete<uint8_t[]>());
        memcpy(mediaInfo.sps.get(), sps.first, mediaInfo.sps_size);

        // pps
        H264Paraser::Nal pps = H264Paraser::findNal(sps.second, extradataSize - (sps.second - (uint8_t*)extradata));
        if (pps.first != nullptr && pps.second != nullptr && (*pps.first & 0x1f) == 8) // pps数据
        {
            mediaInfo.pps_size = pps.second - pps.first + 1;
            mediaInfo.pps.reset(new uint8_t[mediaInfo.pps_size], std::default_delete<uint8_t[]>());
            memcpy(mediaInfo.pps.get(), pps.first, mediaInfo.pps_size);
        }
    }
    uint32_t audioExtraSize = audio_encoder_->GetSpecificConfig(extradata, 1024);

    mediaInfo.audio_specific_config_size = audioExtraSize;
    mediaInfo.audio_specific_config.reset(new uint8_t[mediaInfo.audio_specific_config_size], std::default_delete<uint8_t[]>());
    memcpy(mediaInfo.audio_specific_config.get(), extradata, audioExtraSize);

    pusher_->SetMediaInfo(mediaInfo);
    return true;
}

void RtmpPushManager::Close()
{
    exit_ = true;
    isConnect = false;
    if (pusher_ && pusher_->IsConnected())
    {
        pusher_->Close();
        pusher_.reset();
        pusher_ = nullptr;
    }

    StopEncoder();
    StopCapture();
}

void RtmpPushManager::StopEncoder()
{
    if (audioCaptureThread_)
    {
        audioCaptureThread_->join();
        audioCaptureThread_.reset();
        audioCaptureThread_ = nullptr;
    }

    if (videoCaptureThread_)
    {
        videoCaptureThread_->join();
        videoCaptureThread_.reset();
        videoCaptureThread_ = nullptr;
    }

    if (video_encoder_)
    {
        video_encoder_->Close();
        video_encoder_.reset();
        video_encoder_ = nullptr;
    }

    if (audio_encoder_)
    {
        audio_encoder_->Close();
        audio_encoder_.reset();
        audio_encoder_ = nullptr;
    }
}

void RtmpPushManager::StopCapture()
{
    if (audio_Capture_)
    {
        audio_Capture_->Close();
        audio_Capture_.reset();
        audio_Capture_ = nullptr;
    }

    if (screen_Capture_)
    {
        screen_Capture_->Close();
        screen_Capture_.reset();
        screen_Capture_ = nullptr;
    }
}

// TODO: 搞清楚EncodeVideo 和 EncodeAudio的逻辑
void RtmpPushManager::EncodeVideo()
{
    static Timestamp timestamp;
    uint32_t frameRate = 30;
    while (!exit_ && isConnect)
    {
        uint32_t elapsed = timestamp.Elapsed();
        //获取延迟时间
        uint32_t delay = frameRate;
        if (elapsed > delay)
        {
            delay = 0; //重置延迟
        }
        else
        {
            delay -= elapsed;
        }
        //休眠这个延迟时间
        std::this_thread::sleep_for(std::chrono::milliseconds(delay));
        timestamp.Reset(); //重新获取这个时间
        // FrameContainer bgra_image;
        AVFramePtr bgra_av_frame;
        uint32_t width = 0, height = 0;
        width = screen_Capture_->GetWidth();
        height = screen_Capture_->GetHeight();
        if (screen_Capture_ && video_encoder_ && pusher_)
        {
            if (!screen_Capture_->CaptureFrame(bgra_av_frame)) continue;
            FrameContainer out_frame;
            uint32_t data_size = bgra_av_frame->height * bgra_av_frame->linesize[0];
            if (video_encoder_->Encode(bgra_av_frame->data[0], bgra_av_frame->width, bgra_av_frame->height, data_size, out_frame) > 0)
            {
                if (out_frame.size() > 0) PushVideo(&out_frame[0], out_frame.size());
            }
        }
    }
}

void RtmpPushManager::EncodeAudio()
{
    std::shared_ptr<uint8_t> pcm_buffer(new uint8_t[48000 * 8], std::default_delete<uint8_t[]>());
    // 获取样本数
    uint32_t frameSamples = audio_encoder_->GetFrames();
    while (!exit_ && isConnect) {
        if (audio_Capture_->GetSamples() >= (int)frameSamples) {
            if (audio_Capture_->Read(pcm_buffer.get(), frameSamples) != frameSamples) { continue; }
            AVPacketPtr pkt = audio_encoder_->Encode(pcm_buffer.get(), frameSamples);
            if (pkt) { PushAudio(pkt->data, pkt->size); }
        }
        else
        {
            std::this_thread::sleep_for(std::chrono::milliseconds(1));
        }
    }
}

bool RtmpPushManager::IsKeyFrame(const uint8_t *data, uint32_t size)
{
    int startCode = 0;
    if (data[0] == 0 && data[1] == 0 && data[2] == 0)
    {
        startCode = 3;
    }
    else if (data[0] == 0 && data[1] == 0 && data[2] == 0 && data[3] == 0)
    {
        startCode = 4;
    }

    int type = data[startCode] & 0x1F;
    if (type == 5 || type == 7) // 关键帧
    {
        return true;
    }
    return false;
}

void RtmpPushManager::PushVideo(const quint8 *data, quint32 size)
{
    //准备buffer size = video size - 4 //startcode
    std::shared_ptr<uint8_t> frame(new uint8_t[size - 4], std::default_delete<uint8_t[]>());
    // 拷贝数据
    memcpy(frame.get(), data + 4, size - 4);
    if (size <= 0) return;

    if (pusher_ && pusher_->IsConnected()) {
        pusher_->PushVideoFrame(frame.get(), size - 4);
    }
}

void RtmpPushManager::PushAudio(const quint8 *data, quint32 size)
{
    std::shared_ptr<uint8_t> frame(new uint8_t[size],std::default_delete<uint8_t[]>());
    //拷贝数据
    memcpy(frame.get(),data,size);
    if(size <= 0) return;
    if(pusher_ && pusher_->IsConnected())
    {
        pusher_->PushAudioFrame(frame.get(),size);
    }
}
